package org.example.yuaiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import java.sql.SQLException;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

@Configuration
public class PgVectorVectorStoreConfig {
    @Resource
    private TravelAppDocumentLoader travelAppDocumentLoader;

    @Bean
    public VectorStore pgVectorStore(@Qualifier("postgresJdbcTemplate")JdbcTemplate postgresJdbcTemplate, EmbeddingModel dashscopeEmbeddingModel) throws SQLException {

        float[] embedding = dashscopeEmbeddingModel.embed("测试文本");
        System.out.println("Embedding 维度：" + embedding.length);
        PgVectorStore vectorStore = PgVectorStore.builder(postgresJdbcTemplate, dashscopeEmbeddingModel)
                .dimensions(1536)
                .indexType(HNSW)
                .initializeSchema(true)
                .schemaName("public")
                .vectorTableName("vector_store")
                .maxDocumentBatchSize(10000)
                .build();
        return vectorStore;
    }
}
